import casadi as ca
import numpy as np
import os

def create_flatness_map_functions():
    """
    Create CasADi functions for UAV flatness mapping with automatic differentiation
    and generate C code for forward and backward passes.
    """
    
    # Define symbolic variables matching the C++ implementation
    vel = ca.MX.sym('vel', 3)     # velocity [vx, vy, vz]
    acc = ca.MX.sym('acc', 3)     # acceleration [ax, ay, az]
    jer = ca.MX.sym('jer', 3)     # jerk [jx, jy, jz]
    psi = ca.MX.sym('psi')        # yaw angle
    dpsi = ca.MX.sym('dpsi')      # yaw rate
    
    # Create sub-functions for each step
    
    # 1. Normalize with gradient function
    def normalize_with_grad(x, xd):
        """Normalize vector x and compute derivative of normalized vector"""
        x_sqr_norm = ca.dot(x, x)
        x_norm = ca.sqrt(x_sqr_norm)
        x_nor = x / x_norm
        x_nord = (xd - x * (ca.dot(x, xd) / x_sqr_norm)) / x_norm
        return x_nor, x_nord
    
    # 2. Compute intermediate values
    # zb and zbd from acceleration and jerk
    zb, zbd = normalize_with_grad(acc, jer)
    
    # xc and xcd from yaw angle and rate
    syaw = ca.sin(psi)
    cyaw = ca.cos(psi)
    xc = ca.vertcat(cyaw, syaw, 0.0)
    xcd = ca.vertcat(-syaw * dpsi, cyaw * dpsi, 0.0)
    
    # yc and ycd from cross products
    yc = ca.cross(zb, xc)
    ycd = ca.cross(zbd, xc) + ca.cross(zb, xcd)
    
    # yb and ybd from normalization
    yb, ybd = normalize_with_grad(yc, ycd)
    
    # xb and xbd from cross products
    xb = ca.cross(yb, zb)
    xbd = ca.cross(ybd, zb) + ca.cross(yb, zbd)
    
    # 3. Compute outputs
    # Rotation matrix (using column vectors)
    rotM = ca.horzcat(xb, yb, zb)
    
    # Convert rotation matrix to quaternion
    # Using method from https://en.wikipedia.org/wiki/Rotation_matrix#Quaternion
    tr = rotM[0,0] + rotM[1,1] + rotM[2,2]
    
    # Fix the logical operations using CasADi functions
    cond1 = tr > 0
    cond2 = ca.logic_and(rotM[0,0] > rotM[1,1], rotM[0,0] > rotM[2,2])
    cond3 = rotM[1,1] > rotM[2,2]
    
    w = ca.if_else(cond1,
                   ca.sqrt(1.0 + tr) / 2.0,
                   ca.if_else(cond2,
                              ca.sqrt(1.0 + rotM[0,0] - rotM[1,1] - rotM[2,2]) / 2.0,
                              ca.if_else(cond3,
                                         ca.sqrt(1.0 - rotM[0,0] + rotM[1,1] - rotM[2,2]) / 2.0,
                                         ca.sqrt(1.0 - rotM[0,0] - rotM[1,1] + rotM[2,2]) / 2.0)))
    
    x_q = ca.if_else(cond1,
                     (rotM[2,1] - rotM[1,2]) / (4.0 * w),
                     ca.if_else(cond2,
                                (rotM[0,1] + rotM[1,0]) / (4.0 * w),
                                ca.if_else(cond3,
                                           (rotM[1,2] + rotM[2,1]) / (4.0 * w),
                                           (rotM[2,0] + rotM[0,2]) / (4.0 * w))))
    
    y_q = ca.if_else(cond1,
                     (rotM[0,2] - rotM[2,0]) / (4.0 * w),
                     ca.if_else(cond2,
                                (rotM[2,0] + rotM[0,2]) / (4.0 * w),
                                ca.if_else(cond3,
                                           (rotM[1,0] + rotM[0,1]) / (4.0 * w),
                                           (rotM[0,1] + rotM[1,0]) / (4.0 * w))))
    
    z_q = ca.if_else(cond1,
                     (rotM[1,0] - rotM[0,1]) / (4.0 * w),
                     ca.if_else(cond2,
                                (rotM[1,2] - rotM[2,1]) / (4.0 * w),
                                ca.if_else(cond3,
                                           (rotM[0,2] - rotM[2,0]) / (4.0 * w),
                                           (rotM[2,1] - rotM[1,2]) / (4.0 * w))))
    
    quat_result = ca.vertcat(w, x_q, y_q, z_q)
    
    # Body rates computed from Eq. (109)
    omg_result = ca.vertcat(
        (ca.dot(zb, ybd) - ca.dot(yb, zbd)) / 2.0,
        (ca.dot(xb, zbd) - ca.dot(zb, xbd)) / 2.0,
        (ca.dot(yb, xbd) - ca.dot(xb, ybd)) / 2.0
    )
    
    # Thrust is the norm of the acceleration vector
    thr_result = ca.norm_2(acc)
    
    # Create the main forward function with a valid C identifier name
    forward_func = ca.Function(
        'flatness_forward',  # Use a valid C identifier
        [vel, acc, jer, psi, dpsi],
        [thr_result, quat_result, omg_result],
        ['velocity', 'acceleration', 'jerk', 'yaw', 'yaw_rate'],
        ['thrust', 'quaternion', 'angular_velocity']
    )
    
    # Generate Jacobian (backward pass) for each output
    input_vars = ca.vertcat(vel, acc, jer, psi, dpsi)
    jac_thr = ca.jacobian(thr_result, input_vars)
    jac_quat = ca.jacobian(quat_result, input_vars)
    jac_omg = ca.jacobian(omg_result, input_vars)
    
    # Create backward functions with valid C identifier names
    backward_thrust = ca.Function(
        'flatness_backward_thrust',
        [vel, acc, jer, psi, dpsi],
        [jac_thr],
        ['velocity', 'acceleration', 'jerk', 'yaw', 'yaw_rate'],
        ['jacobian_thrust']
    )
    
    backward_quaternion = ca.Function(
        'flatness_backward_quaternion',
        [vel, acc, jer, psi, dpsi],
        [jac_quat],
        ['velocity', 'acceleration', 'jerk', 'yaw', 'yaw_rate'],
        ['jacobian_quaternion']
    )
    
    backward_omega = ca.Function(
        'flatness_backward_omega',
        [vel, acc, jer, psi, dpsi],
        [jac_omg],
        ['velocity', 'acceleration', 'jerk', 'yaw', 'yaw_rate'],
        ['jacobian_omega']
    )
    
    return forward_func, backward_thrust, backward_quaternion, backward_omega

def generate_c_code():
    """
    Generate C code for the flatness mapping functions
    """
    # Create the output directory if it doesn't exist
    output_dir = "casadi_auto_generated"
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Create functions
    forward_func, backward_thrust, backward_quaternion, backward_omega = create_flatness_map_functions()
    
    # Generate C code for forward function
    forward_opts = {
        'with_header': True,
        'with_mem': True,
        'with_export': True,
        'include_math': True
    }
    
    # Generate files in the output directory
    forward_func.generate(forward_opts)
    # backward_thrust.generate(os.path.join(output_dir, 'flatness_backward_thrust.c'), forward_opts)
    # backward_quaternion.generate(os.path.join(output_dir, 'flatness_backward_quaternion.c'), forward_opts)
    # backward_omega.generate(os.path.join(output_dir, 'flatness_backward_omega.c'), forward_opts)
    
    print("C code generation completed!")
    print("Generated files in directory:", output_dir)
    print("- flatness_forward.c/.h")
    print("- flatness_backward_thrust.c/.h")
    print("- flatness_backward_quaternion.c/.h")
    print("- flatness_backward_omega.c/.h")

def test_functions():
    """
    Test the generated functions with sample data
    """
    # Create functions
    forward_func, backward_thrust, backward_quaternion, backward_omega = create_flatness_map_functions()
    
    # Sample test data
    vel_test = np.array([1.0, 0.5, 0.1])
    acc_test = np.array([0.1, 0.2, 9.81])  # Includes gravity
    jer_test = np.array([0.01, 0.02, 0.1])
    psi_test = 0.5  # radians
    dpsi_test = 0.1  # rad/s
    
    # Evaluate forward function
    [thr, quat, omg] = forward_func(vel_test, acc_test, jer_test, psi_test, dpsi_test)
    
    print("Forward pass results:")
    print(f"Thrust: {float(thr)}")
    print(f"Quaternion: [{float(quat[0])}, {float(quat[1])}, {float(quat[2])}, {float(quat[3])}]")
    print(f"Angular velocity: [{float(omg[0])}, {float(omg[1])}, {float(omg[2])}]")
    
    # Evaluate backward functions
    jac_thr = backward_thrust(vel_test, acc_test, jer_test, psi_test, dpsi_test)
    jac_quat = backward_quaternion(vel_test, acc_test, jer_test, psi_test, dpsi_test)
    jac_omg = backward_omega(vel_test, acc_test, jer_test, psi_test, dpsi_test)
    
    print("\nBackward pass results (Jacobians):")
    print(f"Thrust Jacobian shape: {jac_thr.shape}")
    print(f"Quaternion Jacobian shape: {jac_quat.shape}")
    print(f"Omega Jacobian shape: {jac_omg.shape}")

if __name__ == "__main__":
    # Test the functions
    test_functions()
    
    # Generate C code
    generate_c_code()